# 回测类
import pandas as pd
import copy


class Backtest:

    def __init__(self, Bars, cash, circle, riskfree_rate=0.0001, tax_rate1=0.001, tax_rate2=0.0001, tax_rate3=0.0003):
        """

        :param Bars: list, Bar类的实例化
        :param cash: float, 个人拥有的可用于投资的资金
        :param circle: int, 调仓周期
        :param riskfree_rate: float, 用来计算盈余资金的收益
        :param tax_rate1: float, 印花税千分之一，卖方支付
        :param tax_rate2: float, 过户费万分之0.2，双方支付；其他费用可以加在过户费上面;总的估计为万分之一
        :param tax_rate3: float, 佣金暂定万分之三，双方支付（最小为5）
        """
        self.stock = Bars
        self.cash = cash
        self.circle = circle
        self.position = {}  # 拥有的股票头寸（数量）（手）  每只股票的数量list，对应Bar
        self.changeposition = {}  # 每只股票的操作
        self.stock_values = {}  # 每只股票的价格
        self.stock_weight = {}  # 每只股票占总资金的权重
        self.riskfree_rate = riskfree_rate
        self.tax_rate1 = tax_rate1
        self.tax_rate2 = tax_rate2
        self.tax_rate3 = tax_rate3

    # 买入指令
    def buy(self, price, cash_available, volume):
        price0 = price * 100  # 每手股票总价
        price1 = price0 * (1 + self.tax_rate2 + self.tax_rate3)  # 每手含交易成本的总价
        delta_position0 = cash_available // price1  # 初步判断可买多少手股票
        delta_value0 = delta_position0 * price0  # 初步判断可买股票的价值
        tax2 = delta_value0 * self.tax_rate2
        tax3 = delta_value0 * self.tax_rate3

        # 考虑是否达到最小佣金的情况
        if tax3 < 5:  # 佣金费没有达到最小值
            tax3 = 5
            cash_available0 = (cash_available - tax3) / (1 + self.tax_rate2)  # 可用资金
            delta_position = cash_available0 // price0  # 买入多少手股票
            delta_value = delta_position * price0  # 买入股票的价值
            tax2 = delta_value * self.tax_rate2
        else:  # 佣金达到最小值
            delta_position = delta_position0  # 买入多少手股票
            delta_value = delta_value0  # 买入股票的价值

        if delta_position * 100 > volume:  # 如果购买股票数量大于股票当天成交数量
            delta_position = volume//100
            delta_value = delta_position * price0  # 买入股票的价值
            tax2 = delta_value * self.tax_rate2
            tax3 = delta_value * self.tax_rate3

        total_value = delta_value + tax2 + tax3
        list = [delta_position * 100, total_value]   # [买入股票数量，含交易成本所花总价]
        return list

    # 卖出指令
    def sell(self, position, price, sell_value, volume):
        delta_position = sell_value // price

        if delta_position > position:
            delta_position = position
        elif delta_position > volume:
            print("卖出数量超过当天成交数量")
            delta_position = volume

        stock_value = price * delta_position   # 交易股票价值
        tax1 = stock_value * self.tax_rate1
        tax2 = stock_value * self.tax_rate2
        tax3 = stock_value * self.tax_rate3

        if tax3 < 5:  # 佣金费没有达到最小值
            tax3 = 5

        total_value = stock_value - tax1 - tax2 - tax3
        list = [delta_position, total_value]
        return list

    def get_cash(self):  # 返回可用资金
        return self.cash

    def get_calendar_data(self, df):  # 调用策略，调用buy和sell，调用数据读写计算获得
        """
        回测思路：首先根据调仓周期判断是否为调仓日；
                调仓日：如果是第一天，对停牌股票和非停牌股票所有仓位信息进行初始化；
                    如果不是第一天，非停牌股票根据策略导入的目标权重进行买卖并更新所有仓位信息，停牌股票只需将变化仓位设为0，其余不调整
                非调仓日：停牌股票和非停牌股票分别更新所有仓位信息
                最后，每日末更新现金和股票所占总资金的权重
        :param df: dataFrame, 策略得出的回测期间内的每日目标权重
        :return: dataFrame, 仓位信息日志
        """

        datelist = df.index.tolist()
        record = []  # 用于存储所有日期的股票日志
        n = 0

        for date in datelist:
            list = []  # 用于存储每日的股票日志
            stocklist = []  # 用于存储非停牌股票代码

            for bar in self.stock:
                if bar.stock_data.loc[date, "open"] != float("nan"):  # 非停牌股票
                    stocklist.append(bar.code)
            sumweight = df.loc[date, stocklist].sum()     # 加总所有非停牌股票权重

            if n % self.circle == 0:  # 需要调仓
                if date == datelist[0]:  # 股票仓位初始化
                    for bar in self.stock:
                        if bar.code not in stocklist:  # 停牌股票初始化
                            self.changeposition[bar.code] = 0
                            self.position[bar.code] = 0
                            self.stock_values[bar.code] = 0
                        else:  # 非停牌股票初始化
                            weight = df.loc[date, bar.code] / sumweight  # 按照所有非停牌股票中的相对权重进行购买
                            info = self.buy(bar.stock_data.loc[date, "open"], weight * self.cash,
                                            bar.stock_data.loc[date, "volume"])
                            self.changeposition[bar.code] = info[0]  # 返回股票变化数量
                            self.position[bar.code] = info[0]  # 返回股票现有仓位
                            self.stock_values[bar.code] = self.position[bar.code] * bar.stock_data.loc[
                                date, "close"]  # 返回股票现有价值
                            self.cash = self.cash - info[1]
                            sumweight = sumweight - df.loc[date, bar.code]  # 动态调整相对权重
                else:  # 后续调仓
                    # 处理停牌股票
                    for bar in self.stock:
                        if bar.code not in stocklist:
                            self.changeposition[bar.code] = 0

                    # 处理非停牌股票
                    all_cash = self.cash + sum(
                        [self.stock_values[bar.code] for bar in self.stock if bar.code in stocklist])  # 初始总市值，以上一天为基准
                    inweight = [self.stock_weight[bar.code] for bar in self.stock if bar.code in stocklist]
                    diff = df.loc[date, stocklist] / sumweight - inweight  # 获取非停牌权重的变化
                    diff.sort_values(ascending=True, inplace=True)  # 将非停牌股票权重排序

                    for stockcode in diff.index:
                        for bar in self.stock:
                            if bar.code == stockcode:
                                break  # 定位到stock
                        if diff[bar.code] == 0:
                            self.changeposition[bar.code] = 0
                            self.stock_values[bar.code] = self.position[bar.code] * bar.stock_data.loc[date, "close"]
                        elif diff[bar.code] < 0:  # 权重小于0，就卖
                            info = self.sell(self.position[bar.code], bar.stock_data.loc[date, "open"],
                                             -diff[bar.code] * all_cash, bar.stock_data.loc[date, "volume"])
                            self.changeposition[bar.code] = -info[0]
                            self.position[bar.code] = self.position[bar.code] - info[0]
                            self.stock_values[bar.code] = self.position[bar.code] * bar.stock_data.loc[date, "close"]
                            self.cash = self.cash + info[1]
                        else:  # 权重大于0，买
                            weight = diff[bar.code] / diff[bar.code:].sum()  # 动态调整相对权重
                            info = self.buy(bar.stock_data.loc[date, "open"], weight * self.cash,
                                            bar.stock_data.loc[date, "volume"])
                            self.changeposition[bar.code] = info[0]
                            self.position[bar.code] = self.position[bar.code] + info[0]
                            self.stock_values[bar.code] = self.position[bar.code] * bar.stock_data.loc[date, "close"]
                            self.cash = self.cash - info[1]
            else:  # 不需要调仓
                for bar in self.stock:
                    if bar.code not in stocklist:  # 停牌股票
                        self.changeposition[bar.code] = 0
                    else:  # 非停牌股票
                        self.changeposition[bar.code] = 0
                        self.stock_values[bar.code] = self.position[bar.code] * bar.stock_data.loc[date, "close"]

            self.cash = self.cash * (1 + self.riskfree_rate)  # 更新现金
            for bar in self.stock:  # 每日末按照每日收盘价，更新每只股票占总资产的权重
                self.stock_weight[bar.code] = self.stock_values[bar.code] / (
                            sum(self.stock_values.values()) + self.cash)

            list.append(self.changeposition)
            list.append(self.position)
            list.append(self.cash)
            list.append(self.stock_weight)
            list.append(self.stock_values)
            list.append(sum(self.stock_values.values()))
            list.append(sum(self.stock_values.values()) + self.cash)
            record.append(copy.deepcopy(list))  # 添加每一天的日志
            n = n + 1
        recorddf = pd.DataFrame(record, index=datelist, columns=["操作", "仓位", "现金", "权重", "股票价值",
                                                                 "股票总价值", "资产价值"])
        return recorddf